package com.yifeng.repo.controller.excel.loader.task;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.support.ExcelTypeEnum;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.gomcarter.frameworks.base.exception.CustomException;
import com.google.common.base.Strings;
import com.talkyun.utils.Looper;
import com.yifeng.repo.base.utils.file.FileHelper;
import com.yifeng.repo.controller.excel.loader.biz.BizExportAction;
import com.yifeng.repo.controller.excel.loader.dto.DataTaskConf;
import com.yifeng.repo.controller.excel.loader.dto.DataTaskDto;
import com.yifeng.repo.controller.excel.loader.dto.RowDataDto;
import com.yifeng.repo.controller.excel.loader.redis.ExcelLoaderRedisWorker;
import com.yifeng.repo.controller.excel.loader.task.handler.ExcelFileHandler;
import com.yifeng.repo.controller.excel.loader.task.handler.ExcelWidthStyleStrategy;
import com.yifeng.repo.tookit.oss.OssFileClient;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created by daibing on 2022/12/29.
 */
@Slf4j
public class DataExportTask {
    /**
     * 构造函数传入
     */
    private final DataTaskDto dataTaskDto;
    private final BizExportAction bizExportAction;
    private final DataTaskConf properties;
    private final ExcelLoaderRedisWorker redisWorker;
    private final OssFileClient ossFileClient;

    /**
     * 本地文件：导出excel到本地转存文件
     */
    private File exportDumpFile;
    private File exportDumpTemplateFile;
    private ExcelWriter exportDumpFileWriter;
    private final ConcurrentMap<Integer, WriteSheet> exportDumpFileSheetNo2WriteSheet = new ConcurrentHashMap<>();

    /**
     * 启动阶段初始化
     */
    private BlockingQueue<RowDataDto> rowDataQueue;
    private Looper reader;
    private Looper writer;
    private Looper checker;

    /**
     * 计算器：导出需要偏移量（传递给业务根据偏移量查询数据库）
     */
    private final AtomicInteger readerCount = new AtomicInteger(0);
    private final AtomicInteger writerCount = new AtomicInteger(0);
    private volatile int finishedNumber = 0;
    private volatile String finishedOffsetId = null;
    private volatile int lastSheetNo = 0;
    private volatile boolean takeFinished = false;
    private volatile String dataResultUrl = null;

    public DataExportTask(DataTaskDto dataTaskDto, BizExportAction bizExportAction, DataTaskConf properties,
                          ExcelLoaderRedisWorker redisWorker, OssFileClient ossFileClient) {
        this.dataTaskDto = dataTaskDto;
        this.bizExportAction = bizExportAction;
        this.properties = properties;
        this.redisWorker = redisWorker;
        this.ossFileClient = ossFileClient;
    }

    public void startup() {
        this.initDumpFile();
        this.checkAndUpdateTask();
        this.checkAndDownloadDumpFile();
        this.initDumpFileEasyExcelWrite();
        rowDataQueue = new ArrayBlockingQueue<>(properties.getQueueSize());
        reader = this.buildReader(dataTaskDto);
        writer = this.buildWriter(dataTaskDto);
        checker = this.buildChecker(dataTaskDto);
        reader.start();
        writer.start();
        checker.start();
    }

    public void shutdown() {
        reader.shutdown();
        writer.shutdown();
        checker.shutdown();
    }

    public void shutdown4Error(DataTaskDto dataTaskDto) {
        exportDumpFile.delete();
        exportDumpTemplateFile.delete();
        bizExportAction.failed(dataTaskDto);
        shutdown();
    }

    private void initDumpFile() {
        this.exportDumpFile = new File(dataTaskDto.getDataDumpFilePath());
        this.exportDumpTemplateFile = ExcelFileHandler.getTemplateFile(dataTaskDto.getDataDumpFilePath());
    }

    private void checkAndUpdateTask() {
        // 当前还未开始处理，也就还没有开始偏移
        if (dataTaskDto.getFinishedNumber() == 0) {
            return;
        }
        // 当前已经开始处理了，也就是有偏移了，并且转存文件存在就从偏移量开始处理
        if (this.exportDumpFile.exists()) {
            return;
        }
        // 如果转存文件不存在，就需要从头开始
        dataTaskDto.setFinishedNumber(0);
        dataTaskDto.setFinishedOffsetId("0");
        dataTaskDto.setLastSheetNo(0);
    }

    private void checkAndDownloadDumpFile() {
        if (!this.exportDumpFile.exists()) {
            try {
                // 将模板文件拷贝为转存文件（固定一个线程）
                FileHelper.downloadFile(dataTaskDto.getDataFileUrl(), this.exportDumpTemplateFile, 1);
            } catch (IOException e) {
                throw new CustomException(dataTaskDto.getRequestKey() + " 根据模板文件初始化转存文件失败: " + dataTaskDto.getDataFileUrl() + ", " + exportDumpTemplateFile.getPath());
            }
        }
    }

    private void initDumpFileEasyExcelWrite() {
        this.exportDumpFileWriter = EasyExcel.write()
                .registerWriteHandler(new ExcelWidthStyleStrategy(bizExportAction))
                .withTemplate(exportDumpTemplateFile)
                .file(exportDumpFile).excelType(ExcelTypeEnum.XLSX)
                .build();
    }

    /**
     * 导出读线程主体业务逻辑
     */
    private Looper buildReader(DataTaskDto dataTaskDto) {
        return new Looper(dataTaskDto.getRequestKey() + "-export-reader", 50, 1000) {
            private int sheetNo = dataTaskDto.getLastSheetNo();
            private String extractOffsetId = dataTaskDto.getFinishedOffsetId();
            private boolean insertedEndFlag = false;

            @Override
            protected void loop() throws Throwable {
                // 1. 调度 BizExportAction#source 方法持续获取数据，如果返回空集合就退出读线程
                List<RowDataDto> rowDataDtoList = bizExportAction.source(sheetNo, extractOffsetId, properties.getBatchSize());
                if ((rowDataDtoList.isEmpty() && insertedEndFlag)
                        || (!rowDataDtoList.isEmpty() && rowDataDtoList.get(rowDataDtoList.size() - 1).getRowNumber() == Integer.MIN_VALUE)) {
                    log.info("export task {} reader finished: offsetId={}, readerCount={}, insertedEndFlag={}",
                            dataTaskDto.getRequestKey(), extractOffsetId, readerCount.get(), insertedEndFlag);
                    bizExportAction.readFinished(dataTaskDto);
                    this.shutdown();
                    return;
                }
                if (rowDataDtoList.isEmpty() && !insertedEndFlag) {
                    // 填入一个结束标记位，表示读取excel完成。
                    rowDataDtoList = Collections.singletonList(new RowDataDto(Integer.MIN_VALUE, Collections.emptyMap()));
                    insertedEndFlag = true;
                }

                // 2. 行数据写入内部队列
                for (RowDataDto rowDataDto : rowDataDtoList) {
                    rowDataQueue.put(rowDataDto);
                    if (rowDataDto.getRowNumber() == Integer.MIN_VALUE) {
                        break;
                    }
                    readerCount.getAndIncrement();
                }

                // 3. 更新抽取数据的偏移量，如果是插入的结束标记位，就不需要更新extractOffsetId
                if (insertedEndFlag) {
                    return;
                }
                RowDataDto lastRowDataDto = rowDataDtoList.get(rowDataDtoList.size() - 1);
                String offsetIdFieldName = bizExportAction.offsetIdFieldName(lastRowDataDto.getSheetNo());
                if (!lastRowDataDto.getData().containsKey(offsetIdFieldName)) {
                    throw new CustomException(String.format("抽取业务数据包偏移量失败：requestKey=%s, offsetIdFieldName=%s, sheetNo=%s, rowNumber=%s",
                            dataTaskDto.getRequestKey(), offsetIdFieldName, lastRowDataDto.getSheetNo(), lastRowDataDto.getRowNumber()));
                }
                sheetNo = lastRowDataDto.getSheetNo();
                extractOffsetId = lastRowDataDto.getData().get(offsetIdFieldName);
            }

            @Override
            protected void loopThrowable(Throwable t) {
                log.warn("export task {} reader failed: error=", dataTaskDto.getRequestKey(), t);
                dataTaskDto.setErrorMessage("reader failed by " + t.getMessage());
                shutdown4Error(dataTaskDto);
            }
        };
    }

    /**
     * 导出写线程主体业务逻辑
     */
    private Looper buildWriter(DataTaskDto dataTaskDto) {
        return new Looper(dataTaskDto.getRequestKey() + "-export-writer", 50, 1000) {
            private Integer firstSheetNo;
            private RowDataDto otherSheetRowDataDto;

            @Override
            protected void loop() throws Throwable {
                // 1. 准备一批数据
                List<RowDataDto> rowDataDtoList = new ArrayList<>(properties.getBatchSize() + 1);
                if (otherSheetRowDataDto != null) {
                    firstSheetNo = otherSheetRowDataDto.getSheetNo();
                    rowDataDtoList.add(otherSheetRowDataDto);
                    writerCount.getAndIncrement();
                    otherSheetRowDataDto = null;
                }
                for (int i = 0; i < properties.getBatchSize(); i++) {
                    RowDataDto rowDataDto = rowDataQueue.take();
                    // 读取到结束标记位
                    if (rowDataDto.getRowNumber() == Integer.MIN_VALUE) {
                        takeFinished = true;
                        break;
                    }
                    // 记录本次获取数据的第一个数据包的sheetNo
                    if (firstSheetNo == null) {
                        firstSheetNo = rowDataDto.getSheetNo();
                    }
                    // 换sheet开始写数据
                    if (rowDataDto.getSheetNo() != firstSheetNo) {
                        otherSheetRowDataDto = rowDataDto;
                        break;
                    }
                    rowDataDtoList.add(rowDataDto);
                    writerCount.getAndIncrement();
                }

                // 2. 数据追加模式写入excel转存文件
                if (!rowDataDtoList.isEmpty()) {
                    List<List<String>> rowList = this.buildResultData(rowDataDtoList, bizExportAction.templateHeader(firstSheetNo));
                    exportDumpFileWriter.write(rowList, getExportDumpFileSheet(firstSheetNo));
                    finishedNumber = writerCount.get();
                    RowDataDto lastRowDataDto = rowDataDtoList.get(rowDataDtoList.size() - 1);
                    finishedOffsetId = lastRowDataDto.getData().get(bizExportAction.offsetIdFieldName(firstSheetNo));
                    lastSheetNo = lastRowDataDto.getSheetNo();
                    firstSheetNo = null;
                }

                // 3. 数据处理完成就将转存文件上传到oss
                if (takeFinished) {
                    exportDumpFileWriter.finish();
                    String exportResultFileUri = bizExportAction.getResultFileUri(dataTaskDto.getRequestKey(), dataTaskDto.getType());
                    dataResultUrl = ossFileClient.uploadFile(exportResultFileUri, exportDumpFile.getAbsolutePath());
                    exportDumpFile.delete();
                    exportDumpTemplateFile.delete();
                    log.info("export task {} writer finished: readerCount={}, writerCount={}, dataResultUrl={}",
                            dataTaskDto.getRequestKey(), readerCount.get(), writerCount.get(), dataResultUrl);
                    bizExportAction.writeFinished(dataTaskDto);
                    this.shutdown();
                }
            }

            /**
             * 注意内部的List<String>是一行数据，外面的List是多行数据
             */
            private List<List<String>> buildResultData(List<RowDataDto> rowDataDtoList, Map<Integer, String> templateHeader) {
                List<List<String>> rowList = new ArrayList<>(rowDataDtoList.size());
                for (RowDataDto rowDataDto : rowDataDtoList) {
                    List<String> dataList = new ArrayList<>(templateHeader.size());
                    for (int i = 0; i < templateHeader.size(); i++) {
                        String key = templateHeader.get(i);
                        String value = rowDataDto.getData().get(key);
                        dataList.add(value);
                    }
                    rowList.add(dataList);
                }
                return rowList;
            }

            @Override
            protected void loopThrowable(Throwable t) {
                log.warn("export task {} writer failed: error=", dataTaskDto.getRequestKey(), t);
                dataTaskDto.setErrorMessage("writer failed by " + t.getMessage());
                shutdown4Error(dataTaskDto);
            }
        };
    }

    /**
     * 导出检查线程主体业务逻辑
     */
    private Looper buildChecker(DataTaskDto dataTaskDto) {
        return new Looper(dataTaskDto.getRequestKey() + "-export-checker", 1000, 1000) {
            @Override
            protected void loop() throws Throwable {
                // 1. 更新数据任务对象
                Date now = new Date();
                boolean totalFinished = takeFinished && !Strings.isNullOrEmpty(dataResultUrl);
                dataTaskDto.setDataResultUrl(dataResultUrl);
                dataTaskDto.setTotalFinished(totalFinished);
                dataTaskDto.setTotalNumber(totalFinished ? readerCount.get() : bizExportAction.totalNumber());
                dataTaskDto.setFinishedNumber(finishedNumber);
                dataTaskDto.setFinishedOffsetId(finishedOffsetId);
                dataTaskDto.setLastSheetNo(lastSheetNo);
                dataTaskDto.setTotalFinishedTime(totalFinished ? now : null);
                dataTaskDto.setRefreshTime(now);

                // 2. 写入redis，返回写入内容
                String content = redisWorker.saveDataTask(dataTaskDto);

                // 3. 读线程和写线程都处理完成，并且上传文件到oss得到数据结果url
                if (totalFinished) {
                    log.info("export task {} checker finished: content={}", dataTaskDto.getRequestKey(), content);
                    bizExportAction.totalFinished(dataTaskDto);
                    this.shutdown();
                }
            }

            @Override
            protected void loopThrowable(Throwable t) {
                log.warn("export task {} checker failed: error=", dataTaskDto.getRequestKey(), t);
                dataTaskDto.setErrorMessage("checker failed by " + t.getMessage());
                shutdown4Error(dataTaskDto);
            }
        };
    }

    /**
     * 支持多sheet
     */
    private WriteSheet getExportDumpFileSheet(Integer sheetNo) {
        WriteSheet exportDumpFileSheet = exportDumpFileSheetNo2WriteSheet.get(sheetNo);
        if (exportDumpFileSheet != null) {
            return exportDumpFileSheet;
        }
        synchronized (exportDumpFileSheetNo2WriteSheet) {
            exportDumpFileSheet = exportDumpFileSheetNo2WriteSheet.get(sheetNo);
            if (exportDumpFileSheet != null) {
                return exportDumpFileSheet;
            }
            exportDumpFileSheet = EasyExcel.writerSheet(sheetNo).needHead(false).build();
            exportDumpFileSheetNo2WriteSheet.put(sheetNo, exportDumpFileSheet);
            return exportDumpFileSheet;
        }
    }
}
